{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "k7Lc9I6taO3k"
      },
      "source": [
        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/pinecone-io/examples/blob/master/docs/semantic-search.ipynb) [![Open nbviewer](https://raw.githubusercontent.com/pinecone-io/examples/master/assets/nbviewer-shield.svg)](https://nbviewer.org/github/pinecone-io/examples/blob/master/docs/semantic-search.ipynb)\n",
        "\n",
        "# Semantic Search\n",
        "\n",
        "In this walkthrough, we'll learn how to use Pinecone for semantic search using a multilingual translation dataset. \n",
        "\n",
        "We'll grab English sentences and search over a corpus of related sentences, aiming to find the relevant subset to our query.\n",
        "\n",
        "\n",
        "Semantic search is a form of retrieval that allows you to find documents that are similar in meaning to a given query, irrespective of the words used in each query. \n",
        "\n",
        "Semantic search is often in opposition to lexical search, where keywords are used to identify relevant documents to a given query, though it doesn't have to always be this way!\n",
        "\n",
        " It's super helpful for applications that require an understanding of a query's intent (such as when a user queries with a question over a corpus), or for when traditional lexical search doesn't work (such as in multimodal or multilingual applications).\n",
        "\n",
        "\n",
        "To begin, let's install the following libraries:"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Installation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "q03L1BYEZQfe",
        "outputId": "19ee9678-91b3-4b53-ab2e-f1c6156fa01e"
      },
      "outputs": [],
      "source": [
        "!uv pip install -qU \\\n",
        "  pinecone~=7.3.0 \\\n",
        "  pinecone-notebooks==0.1.1 \\\n",
        "  numpy==2.0.2 \\\n",
        "  datasets==3.5.1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_xmobBcgqKEL"
      },
      "source": [
        "---\n",
        "\n",
        "🚨 _Note: the above `uv pip install` is formatted for Colab Jupyter notebooks. If running elsewhere you may need to drop the `!`._. If you want to run without uv, remove \"uv\"\n",
        "\n",
        "---"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hrSfFiIC5roI"
      },
      "source": [
        "## Setting up"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Get and Set the Pinecone API Key\n",
        "\n",
        "We'll first need a free Pinecone account and API key. \n",
        "\n",
        "This cell will help you create an account if you don't have one and then create an API key and save it in your Colab environment.\n",
        "\n",
        "Run the cell below, and click the Pinecone Connect button to create an account or log in, and follow the prompts to generate an API key:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "✓ Using API key from environment variable\n"
          ]
        }
      ],
      "source": [
        "from pinecone_notebooks.colab import Authenticate\n",
        "\n",
        "Authenticate()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Now that our key is ready, we can retrieve it from our environment and proceed:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mc66NEBAcQHY"
      },
      "outputs": [],
      "source": [
        "from pinecone import Pinecone\n",
        "# Initialize client\n",
        "import os\n",
        "\n",
        "api_key = os.environ.get(\"PINECONE_API_KEY\")\n",
        "\n",
        "pc = Pinecone(\n",
        "        # You can remove this for your own projects!\n",
        "        api_key=api_key,\n",
        "        source_tag=\"pinecone_examples:docs:semantic_search\"\n",
        "    )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Creating a Pinecone Index with Integrated Inference\n",
        "\n",
        "Typically, semantic search requires three pieces: a processed data source (chunks, or records in Pinecone), an embedding model, and a vector database.\n",
        "\n",
        "Integrated Inference allows you to specify the creation of a Pinecone index with a specific Pinecone-hosted embedding model, which makes it easy to interact with the index. To learn more about Integrated Inference, including what other models are available, take a [look here](https://docs.pinecone.io/guides/get-started/overview#integrated-embedding).\n",
        "\n",
        "\n",
        "Here, we specify a starter tier index with the [llama-text-embed-v2](https://docs.pinecone.io/models/llama-text-embed-v2) embedding model. We also specify a mapping for what field in our records we will embed with this model. Then, we grab the index we just created for embedding later.\n",
        "\n",
        "Want to instead embed a subset with multiple languages? Use the [multilingual-e5-large model](https://docs.pinecone.io/models/multilingual-e5-large) and simply specify this inplace of the previous model when creating an index."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "kT8pfoO46Iwg",
        "outputId": "0fec19be-c74d-4602-bec6-24d61cfc5bb4"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{'dimension': 1024,\n",
              " 'index_fullness': 0.0,\n",
              " 'metric': 'cosine',\n",
              " 'namespaces': {},\n",
              " 'total_vector_count': 0,\n",
              " 'vector_type': 'dense'}"
            ]
          },
          "execution_count": 5,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "\n",
        "index_name = \"semantic-search\"\n",
        "\n",
        "if not pc.has_index(index_name):\n",
        "    pc.create_index_for_model(\n",
        "        name=index_name,\n",
        "        cloud=\"aws\",\n",
        "        region=\"us-east-1\",\n",
        "        embed={\n",
        "            # Use this if you want to instead embed non-english or a multilingual subset of the data\n",
        "            #\"model\":\"multilingual-e5-large\",\n",
        "            \"model\": \"llama-text-embed-v2\",\n",
        "            \"field_map\":{\"text\": \"chunk_text\"}\n",
        "        }\n",
        "    )\n",
        "\n",
        "# Initialize index client\n",
        "index = pc.Index(name=index_name)\n",
        "\n",
        "# View index stats\n",
        "index.describe_index_stats()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Creating our dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We're working with a small subset of a large multilingual dataset called Tatoeba. Tatoeba consists of hundreds of thousands of sentence translation pairs, and sometimes serves as a benchmark for crosslingual semantic search capabilities.\n",
        "\n",
        "In this notebook, we're just testing semantic search, so we'll grab a subset of english sentences that include the word \"park\".\n",
        "\n",
        "Why \"park\"? In English, park has multiple meanings which occur in different contexts. It could mean a place, such as a public park. Or, it could mean an action with a car (to park) or a place (park-ing spot). Semantic search using embedding models will naturally distinguise between these contexts, without invervention or labeling!\n",
        "\n",
        "This is the key benefit for semantic search; a way to abstract and represent the meaning of user queries without any additional work. \n",
        "\n",
        "And, since our embedding model is inherently multilingual, we can even do this semantic search across several languages without any additional work!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {},
      "outputs": [],
      "source": [
        "from datasets import load_dataset\n",
        "# specify that we want the english-spanish translation pairs\n",
        "tatoeba = load_dataset(\"Helsinki-NLP/tatoeba\", lang1=\"en\", lang2=\"es\", trust_remote_code=True, split=\"train\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's take a quick look at a few data points"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{'id': ['0', '1', '2', '3', '4'],\n",
              " 'translation': [{'en': \"Let's try something.\", 'es': '¡Intentemos algo!'},\n",
              "  {'en': \"Let's try something.\", 'es': 'Intentemos algo.'},\n",
              "  {'en': \"Let's try something.\", 'es': 'Permíteme hacer algo.'},\n",
              "  {'en': \"Let's try something.\", 'es': 'Permíteme intentarlo.'},\n",
              "  {'en': 'I have to go to sleep.', 'es': 'Tengo que irme a dormir.'}]}"
            ]
          },
          "execution_count": 7,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "tatoeba[0:5]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {},
      "outputs": [],
      "source": [
        "keywords= [\"park\"]\n",
        "\n",
        "def simple_keyword_filter(sentence, keywords):\n",
        "  # filter for a list of keywords by sentence\n",
        "  # This is really just for making a toy example quickly, not useful for production.\n",
        "    for keyword in keywords:\n",
        "        if keyword in sentence:\n",
        "            return True\n",
        "    return False\n",
        "\n",
        "def transform_dataset_for_pinecone(dataset, use_filter=True):\n",
        "    # Feel free to adjust this code to simulate a larger search!\n",
        "\n",
        "    if use_filter:\n",
        "        # filter for a list of keywords by sentence, helpful for building intuition on semantic search\n",
        "        translation_pairs = dataset.filter(lambda x: simple_keyword_filter(\n",
        "        sentence = x[\"translation\"][\"en\"], keywords=keywords))\n",
        "    else:\n",
        "        # use the full 200k+ dataset. Run only if you want to embed this many records!\n",
        "        translation_pairs = dataset\n",
        "\n",
        "    # flatten and shuffle for ease of use\n",
        "    translation_pairs = translation_pairs.flatten()\n",
        "    translation_pairs = translation_pairs.shuffle(seed=1)\n",
        "\n",
        "    # If you want to include the spanish subset, simply repeat the below steps with \"es\" instead of \"en\"\n",
        "    # Be sure to create your index with multilingual-e5-large as well in this case!\n",
        "    english_sentences = translation_pairs.rename_column(\"translation.en\", \"text\").remove_columns(\"translation.es\")\n",
        "\n",
        "    # add lang column to indicate embedding origin\n",
        "    english_sentences = english_sentences.add_column(\"lang\", [\"en\"]*len(english_sentences))\n",
        "\n",
        "\n",
        "    records = []\n",
        "\n",
        "    for idx, sentence in enumerate(english_sentences):\n",
        "        # Here, we create a record for each sentence in the dataset\n",
        "        # The record contains an ID and metadata fields which we can use to filter if desired\n",
        "        # The chunk_text field is the text we will embed\n",
        "        records.append(\n",
        "            {\n",
        "                \"id\": str(idx),\n",
        "                \"chunk_text\": sentence[\"text\"],\n",
        "                \"lang\": sentence[\"lang\"]\n",
        "            }\n",
        "        )\n",
        "\n",
        "    # convert to record format\n",
        "    return records\n",
        "\n",
        "\n",
        "records = transform_dataset_for_pinecone(tatoeba)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Upserting data into the Pinecone index"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Here, we embed and upsert the data into Pinecone. What this means is that each record we formatted above will interact with our embedding model we specified prior, and produce a vector embedding. Then, we take these embedding batches and store them in Pinecone with the additional information we specified, which is also known as metadata. \n",
        "\n",
        "Metadata is handy for things like filtering, like for if you stored several languages in the same index and want to return just one based on metadata. To learn more about metadata, take a [look here](https://docs.pinecone.io/guides/index-data/indexing-overview#metadata).\n",
        "\n",
        "We specify and create a namespace called \"english-sentences\", which is a higher level unit of organization when interacting with Pinecone. \n",
        "\n",
        "Querying on namespaces performs a sort of broad filter to only records that exist in that namespace, which has the nice effect of speeding up searches too. \n",
        "\n",
        "To learn more about namespaces, [look here](https://docs.pinecone.io/guides/index-data/indexing-overview#namespaces)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 49,
          "referenced_widgets": [
            "6c9efda7c7394404814e7466d4a9108c",
            "ee94a2a6c50b4abd89fdea7dc2c51b0b",
            "1af0d3d6e7a0465a91fd27a916a0508a",
            "86cf0f742de14c4f95fa5ddb74e07985",
            "ab1007c06e414b2da35507a90e91e8cb",
            "640ee2f088814e6c80014be268749820",
            "e1d5e0366eee4a96bf4a52ee84698c6a",
            "726dc10996c5403d8b0d661a3a6cdb63",
            "1e213f554e8f4c18a96bc71b56d8582a",
            "c62d67ae4d0e498fbf330db34ca8f9e3",
            "15e5e36e202c40fba34e92adfe8aa8de"
          ]
        },
        "id": "RhR6WOi1huXZ",
        "outputId": "9ae50bfe-3ebf-4d8b-e28e-3230aed0e1d7"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "Upserting records batch: 100%|██████████| 5/5 [00:02<00:00,  1.91it/s]\n"
          ]
        }
      ],
      "source": [
        "from tqdm import tqdm\n",
        "\n",
        "batch_size = 96\n",
        "namespace = \"english-sentences\"\n",
        "\n",
        "\n",
        "# We upsert in batches of 96 to avoid hitting the embedding model's rate limit.\n",
        "# Libraries like backoff can be used here to handler large embedding jobs.\n",
        "\n",
        "for start in tqdm(range(0, len(records), batch_size), f\"Upserting records batch: \"):\n",
        "    index.upsert_records(records=records[start:start+batch_size], namespace = namespace)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VrK_IN079Vuu"
      },
      "source": [
        "## Making Queries"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rr4unPAq9alb"
      },
      "source": [
        "Now that our index is populated we can begin making queries. \n",
        "\n",
        "The tricky part about querying with semantic search is that we'd normally need to involve an embedding model here again too! \n",
        "\n",
        "But with Pinecone's Integrated Inference, we can just invoke our index we created and send the text we want to search with there. Specifically,\n",
        "the search query is vectorized using the same embedding model we specified prior, and we use this vector to find all closest vectors in the database to it to return.\n",
        "\n",
        "Neat!\n",
        "\n",
        "Our goal here is to write a query sentence that uses one form of the word park, and find sentences that use park in a semantically similar manner. So, let's try this:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {},
      "outputs": [],
      "source": [
        "search_query = \"I want to go to the park and relax\"\n",
        "\n",
        "results = index.search(\n",
        "    namespace=namespace,\n",
        "    query={\n",
        "        \"top_k\": 10,\n",
        "        \"inputs\": {\n",
        "            'text': search_query\n",
        "        }\n",
        "    }\n",
        ")\n",
        "\n",
        "for result in results[\"result\"][\"hits\"]:\n",
        "    print(f'Sentence: {result[\"fields\"][\"chunk_text\"]} Semantic Similarity Score: {result[\"_score\"]}\\n')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "And now, let's use the other meaning of the word, park!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {},
      "outputs": [],
      "source": [
        "search_query = \"I need a place to park\"\n",
        "\n",
        "results = index.search(\n",
        "    namespace=namespace,\n",
        "    query={\n",
        "        \"top_k\": 10,\n",
        "        \"inputs\": {\n",
        "            'text': search_query\n",
        "        }\n",
        "    }\n",
        ")\n",
        "\n",
        "for result in results[\"result\"][\"hits\"]:\n",
        "    print(f'Sentence: {result[\"fields\"][\"chunk_text\"]} Semantic Similarity Score: {result[\"_score\"]}\\n')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Wait, how is this working?\n",
        "\n",
        "When performing semantic search with Pinecone's vector database, you are asking the following question: Given this query vector, what are the closest vectors to it in the database?\n",
        "\n",
        "Because of the way embedding models are trained, this closeness in vector space corresponds to similarity in meaning. The exact metric used for our implementation is cosine similarity, which is simply the angle between the input vector and a document vector. For small amounts of vectors, this task is trivial, but what happens when you have hundreds of thousands, millions or even billions? And what about query latency?\n",
        "\n",
        "The magic of Pinecone's vector database is advanced algorithms that can quickly index and do this search on billion-scale vectors effectively!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Demo Cleanup\n",
        "\n",
        "You can go ahead and ask more queries above. When you're done, delete the index to save resources. \n",
        "\n",
        "Congrats, you've just implemented semantic search with Pinecone!\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "id": "-cWdeKzhAtww"
      },
      "outputs": [],
      "source": [
        "pc.delete_index(name=index_name)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2B0zxR6hbf5d"
      },
      "source": [
        "---"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "pinecone-examples",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.11.9"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 4
}
